#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
/* 函数原型增加 weight */
extern void binary_cross_entropy_grad_custom_do(uint32_t blockDim, void *stream,
                                                uint8_t *p, uint8_t *y, uint8_t *grad,
                                                uint8_t *weight,   // 新增
                                                uint8_t *z);
#else
#include "tikicpulib.h"
/* 核函数原型增加 weight */
extern "C" __global__ __aicore__
void binary_cross_entropy_grad_custom(GM_ADDR p, GM_ADDR y, GM_ADDR grad,
                                      GM_ADDR weight, GM_ADDR z);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t ByteSize = 8 * 2048 * sizeof(uint32_t);

#ifdef ASCENDC_CPU_DEBUG
    /* 1. 为 weight 多申请一块 Gm 内存 */
    uint8_t *p     = (uint8_t *)AscendC::GmAlloc(ByteSize);
    uint8_t *y     = (uint8_t *)AscendC::GmAlloc(ByteSize);
    uint8_t *grad  = (uint8_t *)AscendC::GmAlloc(ByteSize);
    uint8_t *weight= (uint8_t *)AscendC::GmAlloc(ByteSize);   // 新增
    uint8_t *z     = (uint8_t *)AscendC::GmAlloc(ByteSize);

    ReadFile("./input/input_p.bin",     ByteSize, p,     ByteSize);
    ReadFile("./input/input_y.bin",     ByteSize, y,     ByteSize);
    ReadFile("./input/input_grad.bin",  ByteSize, grad,  ByteSize);
    ReadFile("./input/weight.bin",      ByteSize, weight,ByteSize); // 新增

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    /* 2. 调用参数增加 weight */
    ICPU_RUN_KF(binary_cross_entropy_grad_custom, blockDim,
                p, y, grad, weight, z);

    WriteFile("./output/output_z.bin", z, ByteSize);

    AscendC::GmFree((void *)p);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)grad);
    AscendC::GmFree((void *)weight);   // 新增
    AscendC::GmFree((void *)z);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    /* 3. 主机 / 设备指针各增加 weight */
    uint8_t *pHost, *yHost, *gradHost, *weightHost, *zHost;
    uint8_t *pDevice, *yDevice, *gradDevice, *weightDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&pHost),     ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost),     ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&gradHost),  ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&weightHost),ByteSize)); // 新增
    CHECK_ACL(aclrtMallocHost((void **)(&zHost),     ByteSize));

    CHECK_ACL(aclrtMalloc((void **)&pDevice,     ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice,     ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&gradDevice,  ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&weightDevice,ByteSize, ACL_MEM_MALLOC_HUGE_FIRST)); // 新增
    CHECK_ACL(aclrtMalloc((void **)&zDevice,     ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_p.bin",     ByteSize, pHost,     ByteSize);
    ReadFile("./input/input_y.bin",     ByteSize, yHost,     ByteSize);
    ReadFile("./input/input_grad.bin",  ByteSize, gradHost,  ByteSize);
    ReadFile("./input/weight.bin",      ByteSize, weightHost,ByteSize); // 新增

    CHECK_ACL(aclrtMemcpy(pDevice,     ByteSize, pHost,     ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice,     ByteSize, yHost,     ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(gradDevice,  ByteSize, gradHost,  ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(weightDevice,ByteSize, weightHost,ByteSize, ACL_MEMCPY_HOST_TO_DEVICE)); // 新增

    /* 4. 调用增加 weight */
    binary_cross_entropy_grad_custom_do(blockDim, stream,
                                        pDevice, yDevice, gradDevice,
                                        weightDevice,   // 新增
                                        zDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, ByteSize, zDevice, ByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, ByteSize);

    /* 5. 释放新增资源 */
    CHECK_ACL(aclrtFree(pDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(gradDevice));
    CHECK_ACL(aclrtFree(weightDevice));   // 新增
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(pHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(gradHost));
    CHECK_ACL(aclrtFreeHost(weightHost)); // 新增
    CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}